type Token = ['num', number] | ['sym', string] | ['nm', string] | ['eof', 'EOF']

type Expr =
  | number
  | string
  | ['fn', string, Expr[]]
  | ['un', string, Expr]
  | ['bin', string, Expr, Expr]
  | ['chain', Expr, [string, Expr][]]
  | ['cond', Expr, Expr, Expr]

const INFIX = {
  '?': [10, 11],

  '->': [21, 20],

  '||': [30, 31],

  '&&': [40, 41],

  '=': [60, 61],
  '==': [60, 61],
  '<': [60, 61],
  '>': [60, 61],
  '<=': [60, 61],
  '>=': [60, 61],

  '+': [100, 101],
  '-': [100, 101],

  '*': [110, 111],
  '/': [110, 111],
  '%': [110, 111],

  '^': [131, 130],
  '**': [131, 130],
} satisfies Record<string, [number, number]>

const PREFIX = {
  '!': 50,

  '+': 120,
  '-': 120,
} satisfies Record<string, number>

const LEXER = new RegExp(`(?:${[
  ['num', '[0-9]+(?:\\.[0-9]+)?(?:[eE][+-]?[0-9]+)?'],
  ['nm', '\\p{XIDS}\\p{XIDC}*'],
  ['sym', [...'(,:)', ...Object.keys(INFIX), ...Object.keys(PREFIX)]
    .sort((a, b) => b.length - a.length)
    .map(x => x.replace(/([()|/+*^?])/g, '\\$1')).join('|')],
  ['eof', '$'],
  ['err', '.*']
].map(([g, r]) => `(?<${g}>${r})`).join('|')})\\s*`, 'uy')

function lex(s: string) {
  const out: Token[] = []
  s = s.trimStart()
  LEXER.lastIndex = 0
  for (; ;) {
    const res = LEXER.exec(s)!.groups!
    if (res.num != null) {
      out.push(['num', Number(res.num)])
    } else if (res.nm != null) {
      out.push(['nm', res.nm])
    } else if (res.sym != null) {
      out.push(['sym', res.sym])
    } else if (res.eof != null) {
      return out
    } else {
      throw `无法解析："${res.err}"`
    }
  }
}

function isKey<T>(obj: T, key: PropertyKey): key is keyof T {
  return Object.prototype.hasOwnProperty.call(obj, key)
}

// pratt parser
function parse(t: Token[]) {
  let i = 0
  function peek() {
    return t[i] ?? ['eof', 'EOF']
  }
  function eat() {
    const nxt = peek()
    i++
    return nxt
  }
  function isSym<T extends string>(tok: Token, sym: T): tok is ['sym', T] {
    return tok[0] == 'sym' && tok[1] == sym
  }

  const e = rec(0)
  if (peek()[0] != 'eof') throw `应结束而非 "${peek()[1]}"`
  return e

  function rec(minP: number): Expr {
    let tok = eat()
    let lhs: Expr
    if (tok[0] == 'eof') throw '不应结束'
    if (tok[0] != 'sym') {
      lhs = tok[1]
    } else if (tok[1] == '(') {
      lhs = rec(0)
      tok = eat()
      if (!isSym(tok, ')')) throw `应为右括号而非 "${tok[1]}"`
    } else {
      if (!isKey(PREFIX, tok[1])) throw `"${tok[1]}" 不是前缀运算`
      const rhs = rec(PREFIX[tok[1]])
      lhs = ['un', tok[1], rhs]
    }

    let chaining = false
    for (; ;) {
      tok = peek()
      if (tok[0] == 'eof') break

      if (isSym(tok, '(')) {
        if (typeof lhs != 'string') throw `应为函数而非 ${toString(lhs)}`
        eat()
        const args: Expr[] = []
        if (isSym(peek(), ')')) eat()
        else for (; ;) {
          args.push(rec(0))
          tok = eat()
          if (isSym(tok, ')')) break
          if (!isSym(tok, ',')) throw `应为逗号或右括号而非 "${tok[1]}"`
        }
        lhs = ['fn', lhs, args]
        chaining = false
        continue
      }

      if (tok[0] != 'sym' || !isKey(INFIX, tok[1])) break
      const [lP, rP] = INFIX[tok[1]]
      if (lP < minP) break
      eat()

      if (tok[1] == '?') {
        const mhs = rec(0)
        tok = eat()
        if (!isSym(tok, ':')) throw `应为冒号而非 "${tok[1]}"`
        lhs = ['cond', lhs, mhs, rec(rP)]
        chaining = false
      } else {
        const rhs = rec(rP)
        if (!['=', '==', '<', '>', '<=', '>='].includes(tok[1])) {
          lhs = ['bin', tok[1], lhs, rhs]
          chaining = false
        } else if (chaining && Array.isArray(lhs) && lhs[0] === 'chain') {
          lhs[2].push([tok[1], rhs])
        } else {
          lhs = ['chain', lhs, [[tok[1], rhs]]]
          chaining = true
        }
      }
    }
    return lhs
  }
}

function toString(e: Expr): string {
  return !Array.isArray(e) ? String(e) :
    e[0] == 'fn' ? `${e[1]}(${e[2].map(toString).join(', ')})` :
      e[0] == 'un' ? `(${e[1]}${toString(e[2])})` :
        e[0] == 'bin' ? `(${toString(e[2])} ${e[1]} ${toString(e[3])})` :
          e[0] == 'cond' ? `(${toString(e[1])} ? ${toString(e[2])} : ${toString(e[3])})` :
            `(${toString(e[1])} ${e[2].map(([s, rhs]) => ` ${s} ${rhs}`).join('')}`
}

const CONST = {
  NaN: NaN,
  Infinity: Infinity,
  E: Math.E,
  PI: Math.PI,
} satisfies Partial<Record<string, number>>

const FUNC = {
  abs: 1, acos: 1, acosh: 1, asin: 1, asinh: 1, atan: 1, atanh: 1, cbrt: 1, ceil: 1, cos: 1, cosh: 1, exp: 1, expm1: 1, floor: 1, fround: 1, log: 1, log1p: 1, log10: 1, log2: 1, round: 1, sign: 1, sin: 1, sinh: 1, sqrt: 1, tan: 1, tanh: 1, trunc: 1,

  atan2: 2, pow: 2,

  hypot: -1, max: -1, min: -1
} satisfies Partial<Record<keyof Math, number>>

type Type = 'bool' | 'num'

function check(e: Expr) {
  function unify(e: Expr, expe: Type, get: Type) {
    if (expe != get) throw `${toString(e)} 应为 ${expe} 而非 ${get}`
  }
  unify(e, 'bool', rec(e))
  function rec(e: Expr): Type {
    if (!Array.isArray(e)) return 'num'
    switch (e[0]) {
      case 'fn': {
        e[2].forEach(e => unify(e, 'num', rec(e)))
        return 'num'
      }
      case 'un': {
        const ty = e[1] == '!' ? 'bool' : 'num'
        unify(e[2], ty, rec(e[2]))
        return ty
      }
      case 'bin': {
        const logicOp = ['&&', '||', '->'].includes(e[1])
        const lhs = rec(e[2])
        if (logicOp) unify(e[2], 'bool', lhs)
        unify(e[3], lhs, rec(e[3]))
        return logicOp ? 'bool' : 'num'
      }
      case 'chain': {
        let lhs = e[1], lhsTy = rec(lhs)
        for (const [s, rhs] of e[2]) {
          const rhsTy = rec(rhs)
          if (s != '=' && s != '==') unify(lhs, lhsTy, 'num')
          unify(rhs, lhsTy, rhsTy)
          lhs = rhs
          lhsTy = rhsTy
        }
        return 'bool'
      }
      case 'cond': {
        unify(e[1], 'bool', rec(e[1]))
        const lhs = rec(e[2])
        unify(e[3], lhs, rec(e[3]))
        return lhs
      }
    }
  }
}

function generate(e: Expr, free: Record<string, 1>): string {
  if (typeof e == 'number') return String(e)
  if (typeof e == 'string') {
    if (isKey(CONST, e)) return String(CONST[e])
    free[e] = 1
    return e
  }
  switch (e[0]) {
    case 'fn': {
      if (!isKey(FUNC, e[1])) throw `"${e[1]}" 不是函数`
      if (FUNC[e[1]] != -1 && FUNC[e[1]] != e[2].length)
        throw `需要 ${FUNC[e[1]]} 个参数，但给了 ${e[2].length} 个`
      const args = e[2].map(e => generate(e, free)).join(', ')
      return `Math.${e[1]}(${args})`
    }
    case 'un': return `(${e[1]}${generate(e[2], free)})`
    case 'bin': {
      const lhs = generate(e[2], free)
      const rhs = generate(e[3], free)
      switch (e[1]) {
        case '^': return `(${lhs} ** ${rhs})`
        case '->': return `(!${lhs} || ${rhs})`
        default: return `(${lhs} ${e[1]} ${rhs})`
      }
    }
    case 'chain': {
      const term = (s: string, lhs: string, rhs: string) => {
        switch (s) {
          case '=': return `Object.is(${lhs}, ${rhs})`
          case '==': return `(${lhs} === ${rhs})`
          default: return `(${lhs} ${s} ${rhs})`
        }
      }
      return e[2].length == 1
        ? term(e[2][0][0], generate(e[1], free), generate(e[2][0][1], free))
        : `(($0, ${e[2].map((_, i) => `$${i + 1}`).join(', ')}) => ${e[2].map(([s], i) => term(s, `$${i}`, `$${i + 1}`)).join(' && ')})(${generate(e[1], free)}, ${e[2].map(([, v]) => generate(v, free)).join(', ')})`
    }
    case 'cond': return `(${generate(e[1], free)
      } ? ${generate(e[2], free)
      } : ${generate(e[3], free)})`
  }
}

export function compile(s: string) {
  try {
    const free: Record<string, 1> = {}
    const expr = parse(lex(s))
    check(expr)
    const out = generate(expr, free)
    return { kind: 'ok', expr: out, args: Object.keys(free) } as const
  } catch (e) {
    return { kind: 'err', e } as const
  }
}
